# compare_wildrgbd.py

import os
import os.path as osp
import sys
import glob
import logging
import argparse
import numpy as np
import cv2
from scipy.spatial.transform import Rotation

# Try to import scikit-image for SSIM comparison
try:
    from skimage.metrics import structural_similarity
    ssim_available = True
except ImportError:
    ssim_available = False
    print("Warning: scikit-image not found. SSIM comparison will be skipped.")
    print("         Install it with: pip install scikit-image")


def setup_logging(log_file='logs/comparison.log'):
    """Configures logging to write to a file and the console."""
    log_dir = osp.dirname(log_file)
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
        
    logger = logging.getLogger('ComparisonLogger')
    logger.setLevel(logging.INFO)
    if logger.hasHandlers():
        logger.handlers.clear()
    
    file_handler = logging.FileHandler(log_file, mode='w')
    file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(file_formatter)
    logger.addHandler(file_handler)
    
    console_handler = logging.StreamHandler()
    console_formatter = logging.Formatter('%(message)s')
    console_handler.setFormatter(console_formatter)
    logger.addHandler(console_handler)
    return logger

def compare_images(path1, path2):
    """Compares two images using MAE and SSIM."""
    img1 = cv2.imread(path1)
    img2 = cv2.imread(path2)
    print(img1.shape, img2.shape, img1.max(), img1.mean(), img2.max(), img2.mean())    

    if img1 is None or img2 is None:
        return {"error": "Could not read one or both images."}
    assert img1.shape == img2.shape, {"error": f"Shape mismatch: {img1.shape} vs {img2.shape}"}

    mae = np.mean(np.abs(img1.astype(np.float32) - img2.astype(np.float32)))
    mae_rel = np.mean(np.abs(img1.astype(np.float32) - img2.astype(np.float32))[img1>1e-3]/img1[img1>1e-3])
    
    ssim = None
    if ssim_available:
        # multichannel=True is deprecated and will be removed in scikit-image 1.0. 
        # Use channel_axis=-1 instead.
        try:
            ssim = structural_similarity(img1, img2, channel_axis=-1, data_range=255)
        except TypeError: # Fallback for older scikit-image versions
             ssim = structural_similarity(img1, img2, multichannel=True, data_range=255)


    return {"mae": mae, "mae_rel": mae_rel, "ssim": ssim}

def compare_depths(path1, path2):
    """Compares two depth maps using MAE and RMSE on valid pixels."""
    depth1 = cv2.imread(path1, cv2.IMREAD_UNCHANGED)
    depth2 = cv2.imread(path2, cv2.IMREAD_UNCHANGED)

    if depth1 is None or depth2 is None:
        return {"error": "Could not read one or both depth maps."}
    if depth1.shape != depth2.shape:
        return {"error": f"Shape mismatch: {depth1.shape} vs {depth2.shape}"}

    depth1 = depth1.astype(np.float32)
    depth2 = depth2.astype(np.float32)

    # Create a mask of valid, overlapping pixels (where both are > 0)
    valid_mask = (depth1 > 0) & (depth2 > 0)
    
    if np.sum(valid_mask) == 0:
        return {"mae": 0, "rmse": 0, "valid_pixels": 0}

    diff = np.abs(depth1[valid_mask] - depth2[valid_mask])
    mae = np.mean(diff)
    rmse = np.sqrt(np.mean(diff**2))

    return {"mae": mae, "rmse": rmse, "valid_pixels": int(np.sum(valid_mask))}

def compare_cameras(path1, path2):
    """Compares two camera .npz files for intrinsics and pose differences."""
    try:
        cam1 = np.load(path1)
        cam2 = np.load(path2)
    except Exception as e:
        return {"error": f"Could not load .npz files: {e}"}

    # Intrinsics comparison
    K1, K2 = cam1['intrinsics'], cam2['intrinsics']
    intrinsics_diff = {
        "fx_diff": abs(K1[0, 0] - K2[0, 0]),
        "fy_diff": abs(K1[1, 1] - K2[1, 1]),
        "cx_diff": abs(K1[0, 2] - K2[0, 2]),
        "cy_diff": abs(K1[1, 2] - K2[1, 2]),
    }

    # Pose comparison
    pose1, pose2 = cam1['pose'], cam2['pose']
    
    # Translation difference (Euclidean distance)
    t1, t2 = pose1[:3, 3], pose2[:3, 3]
    translation_dist = np.linalg.norm(t1 - t2)

    # Rotation difference (angular distance in degrees)
    R1, R2 = pose1[:3, :3], pose2[:3, :3]
    # The difference rotation is R_diff = R2 @ R1.T
    # The angle of this rotation is the angular distance between R1 and R2
    R_diff = R2 @ R1.T
    # Use scipy to safely get the angle from the rotation matrix
    angle_rad = Rotation.from_matrix(R_diff).as_rotvec()
    rotation_angle_deg = np.linalg.norm(angle_rad) * 180.0 / np.pi

    return {
        "intrinsics_diff": intrinsics_diff,
        "translation_dist": translation_dist,
        "rotation_angle_deg": rotation_angle_deg
    }


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Compare two versions of a preprocessed wildrgbd dataset."
    )
    parser.add_argument(
        "--dir1", 
        type=str,
        default='/lc/data/3D/wildrgbd/scalerotate/apple/scene_000',
        help="Path to the first processed directory."
    )
    parser.add_argument(
        "--dir2", 
        type=str,
        default='/lc/data/3D/wildrgbd/rotatescale/apple/scene_000',
        help="Path to the second processed directory."
    )
    parser.add_argument(
        "--log_file", 
        type=str, 
        default="logs/comparison_wildrgbd.log",
        help="Path to the output log file."
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print detailed comparison results for every file."
    )
    args = parser.parse_args()
    
    logger = setup_logging(args.log_file)

    logger.info(f"Comparing directories:\n  DIR 1: {args.dir1}\n  DIR 2: {args.dir2}")

    # Find common basenames to compare
    basenames1 = {f.replace('_rgb.png', '') for f in os.listdir(args.dir1) if f.endswith('_rgb.png')}
    basenames2 = {f.replace('_rgb.png', '') for f in os.listdir(args.dir2) if f.endswith('_rgb.png')}
    common_basenames = sorted(list(basenames1.intersection(basenames2)))

    if not common_basenames:
        logger.error("No common files found to compare. Check your directories and file naming.")
        sys.exit(1)

    logger.info(f"Found {len(common_basenames)} common file sets to compare.")
    if len(basenames1) != len(basenames2):
        logger.warning(f"Directories have different file counts ({len(basenames1)} vs {len(basenames2)}).")
        logger.warning(f"Missing in DIR 2: {basenames1 - basenames2}")
        logger.warning(f"Extra in DIR 2: {basenames2 - basenames1}")

    all_results = []

    for basename in common_basenames:
        if args.verbose:
            logger.info(f"\n--- Comparing: {basename} ---")
        
        results = {"basename": basename}

        # Compare RGB
        rgb1_path = osp.join(args.dir1, f"{basename}_rgb.png")
        rgb2_path = osp.join(args.dir2, f"{basename}_rgb.png")
        results["rgb"] = compare_images(rgb1_path, rgb2_path)
        if args.verbose: logger.info(f"  RGB: {results['rgb']}")

        # Compare Depth
        depth1_path = osp.join(args.dir1, f"{basename}_depth.png")
        depth2_path = osp.join(args.dir2, f"{basename}_depth.png")
        results["depth"] = compare_depths(depth1_path, depth2_path)
        if args.verbose: logger.info(f"  Depth: {results['depth']}")

        # Compare Mask
        mask1_path = osp.join(args.dir1, f"{basename}_mask.png")
        mask2_path = osp.join(args.dir2, f"{basename}_mask.png")
        results["mask"] = compare_images(mask1_path, mask2_path)
        if args.verbose: logger.info(f"  Mask: {results['mask']}")

        # Compare Camera
        cam1_path = osp.join(args.dir1, f"{basename}_cam.npz")
        cam2_path = osp.join(args.dir2, f"{basename}_cam.npz")
        results["camera"] = compare_cameras(cam1_path, cam2_path)
        if args.verbose: logger.info(f"  Camera: {results['camera']}")

        all_results.append(results)

    # --- Generate Summary Report ---
    logger.info("\n" + "="*50)
    logger.info("--- COMPARISON SUMMARY ---")
    logger.info("="*50)

    # Helper to calculate averages, ignoring errors
    def avg(key_path):
        values = []
        for r in all_results:
            val = r
            try:
                for k in key_path.split('.'): val = val[k]
                values.append(val)
            except (KeyError, TypeError):
                continue
        return np.mean(values) if values else 0

    logger.info(f"[RGB Image]\n  Avg. MAE: {avg('rgb.mae'):.4f}\n Avg. MAE_Rel: {avg('rgb.mae_rel'):.4f}\n Avg. SSIM: {avg('rgb.ssim'):.4f}")
    logger.info(f"[Mask Image]\n  Avg. MAE: {avg('mask.mae'):.4f}\n  Avg. SSIM: {avg('mask.ssim'):.4f}")
    logger.info(f"[Depth Map]\n  Avg. MAE: {avg('depth.mae'):.4f}\n  Avg. RMSE: {avg('depth.rmse'):.4f}")
    logger.info(f"[Camera Pose]\n  Avg. Translation Distance: {avg('camera.translation_dist'):.4f}\n  Avg. Rotation Angle (deg): {avg('camera.rotation_angle_deg'):.4f}")
    logger.info(f"[Camera Intrinsics]\n  Avg. fx diff: {avg('camera.intrinsics_diff.fx_diff'):.4f}\n  Avg. fy diff: {avg('camera.intrinsics_diff.fy_diff'):.4f}\n  Avg. cx diff: {avg('camera.intrinsics_diff.cx_diff'):.4f}\n  Avg. cy diff: {avg('camera.intrinsics_diff.cy_diff'):.4f}")
    logger.info("="*50)

